from langchain_community.chat_message_histories import SQLChatMessageHistory
from langchain_core.chat_history import InMemoryChatMessageHistory
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnableWithMessageHistory

from langchain_demo.my_llm import llm

# 1、提示词模板
prompt = ChatPromptTemplate.from_messages([
    ('system', '你是一个乐于助人的助手。尽你所能回答所有问题。提供的聊天历史包含与你对话用户的相关信息。'),
    MessagesPlaceholder(variable_name='chat_history', optional=True),
    ('human', '{input}')
])

chain = prompt | llm

# 2、存储聊天记录
store = {}


def get_session_history(session_id: str):
    """从sqlite数据库中的历史消息列表中 返回当前会话的历史消息"""
    return SQLChatMessageHistory(
        session_id=session_id,
        connection_string='sqlite:///chat_history.db',
    )

# langchain所有消息类型 ： SystemMessage, HumanMessage, AIMessage, ToolMessage

# 3、创建带历史记录的处理链
chain_with_message_history = RunnableWithMessageHistory(
    chain,
    get_session_history,
    input_messages_key='input',
    history_messages_key='chat_history',
)

# 4、剪辑 摘要上下文，历史记录。
# 保留最近的前2条消息，把之前的消息形成摘要

# result1 = chain_with_message_history.invoke({'input': '你好，我是张学良'}, config={"configurable": {"session_id": "user123"}})
# print(result1)
#
# result2 = chain_with_message_history.invoke({'input': '你好，历史上有多少人跟我同名'}, config={"configurable": {"session_id": "user123"}})
# print(result2)

result3 = chain_with_message_history.invoke({'input': '你好，你知道我的母亲是谁吗'}, config={"configurable": {"session_id": "user123"}})
print(result3)

